package com.bxmdm.ragdemo.controller;


import com.bxmdm.ragdemo.service.DocumentService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Media;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;
import java.util.List;

@RestController
@RequestMapping("/api/chat")
public class ChatController {

    @Resource
    private OllamaChatClient ollamaChatClient;


    @RequestMapping("/ollama/msg")
    public Object msg(@RequestParam("msg") String msg) {
        ChatResponse response = ollamaChatClient.call(
                new Prompt(msg )
        );
        return response.getResult().getOutput();
    }

    /**
     * 使用7b
     * @param msg
     * @return
     */
    @RequestMapping("/ollama/msg7b")
    public Object msg7b(@RequestParam("msg") String msg) {
        ChatResponse response = ollamaChatClient.call(
                new Prompt(msg,
                        OllamaOptions.create()
                                .withModel("qwen:7b")
                )
        );
        return response.getResult().getOutput();
    }


    @GetMapping("/ai/test")
    public void test() throws IOException {
        byte[] imageData = new ClassPathResource("/multimodal.test.png").getContentAsByteArray();
        Media media = new Media(MimeTypeUtils.IMAGE_PNG, imageData);
        UserMessage userMessage = new UserMessage("简单描述一下这个图片",
                List.of(media));

        ChatResponse response = ollamaChatClient.call(
                new Prompt(List.of(userMessage),
                        OllamaOptions.create().withModel("qwen:7b").withMainGPU(1)
                )
        );
        System.out.println(response.getResult().getOutput().getContent());
    }

    @GetMapping("/ai/mp3")
    public void mp3() throws IOException {
        byte[] imageData = new ClassPathResource("/output.mp3").getContentAsByteArray();
        Media media = new Media(MimeTypeUtils.IMAGE_PNG, imageData);
        UserMessage userMessage = new UserMessage("这个音频文件说了什么",
                List.of(media));

        ChatResponse response = ollamaChatClient.call(
                new Prompt(List.of(userMessage),
                        OllamaOptions.create().withModel("qwen:7b").withMainGPU(1)
                )
        );
        System.out.println(response.getResult().getOutput().getContent());
    }


    @GetMapping("/ai/test1")
    public void test1() throws IOException {
        byte[] imageData = new ClassPathResource("/img.png").getContentAsByteArray();
        Media media = new Media(MimeTypeUtils.IMAGE_PNG, imageData);
        UserMessage userMessage = new UserMessage("我给你的这个图片你帮我描述一下 ?",
                List.of(media));

        ChatResponse response = ollamaChatClient.call(
                new Prompt(List.of(userMessage),
                        OllamaOptions.create().withModel("llava:7b").withMainGPU(1)
                )
        );
        System.out.println(response.getResult().getOutput().getContent());
    }

}
